# 加载环境变量
from openai import OpenAI
from dotenv import load_dotenv, find_dotenv
import os
import json
import mysql.connector

_ = load_dotenv(find_dotenv())  # 读取本地 .env 文件，里面定义了 OPENAI_API_KEY

client = OpenAI(
    api_key=os.getenv("OPENAI_API_KEY"),
    base_url=os.getenv("OPENAI_BASE_URL")
)


# 移除字典中的 None 值
def remove_none_values(data):
    if isinstance(data, dict):
        return {key: remove_none_values(value) for key, value in data.items() if value is not None}
    elif isinstance(data, list):
        return [remove_none_values(item) for item in data if item is not None]
    else:
        return data


# 打印参数。如果参数是有结构的（如字典或列表），则以格式化的 JSON 形式打印；
def print_json(data):
    if (isinstance(data, (list, dict))):
        print(json.dumps(data, indent=4, ensure_ascii=False))
    else:
        print(data)


# 对象序列化
def serialize_json(data):
    if hasattr(data, 'model_dump_json'):
        data = json.loads(data.model_dump_json())

    if (isinstance(data, (list, dict))):
        data = remove_none_values(data)
    return data


def get_sql_completion(params, model="gpt-4"):
    print("=====GPT请求=====")
    messages = serialize_json(params)
    print_json(messages)

    response = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=0,
        seed=1024,  # 随机种子保持不变，temperature 和 prompt 不变的情况下，输出就会不变
        tool_choice="auto",  # 默认值，由 GPT 自主决定返回 function call 还是返回文字回复。也可以强制要求必须调用指定的函数，详见官方文档
        tools=[{
            "type": "function",
            "function": {
                "name": "get_table_ddl",
                "description": "根据数据库表名，获取表的DDL",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "table_names": {
                            "type": "string",
                            "description": "数据库表名，必须是英文，表名中可能包含下划线。多个表名之间用英文逗号分隔。",
                        }
                    },
                    "required": ["table_names"],
                }
            }
        },
        {
            # 摘自 OpenAI 官方示例 https://github.com/openai/openai-cookbook/blob/main/examples/How_to_call_functions_with_chat_models.ipynb
            "type": "function",
            "function": {
                "name": "ask_database",
                "description": "Use this function to answer user questions about business. \
                                Output should be a fully formed SQL query.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": f"""
                                SQL query extracting info to answer the user's question.
                                SQL should be written using this database schema:
                                {database_schema_string}
                                The query should be returned in plain text, not in JSON.
                                The query should only contain grammars supported by SQLite.
                                """,
                        }
                    },
                    "required": ["query"],
                }
            }
        }],
    )
    result = serialize_json(response.choices[0].message)
    print("=====GPT回复=====")
    print_json(result)
    return result


#  描述数据库表结构
database_schema_string = ""


# 获取表信息
def get_table_ddl(table_names):
    name_list = table_names.split(",")
    global cursor
    db_config = {
        'host': os.getenv("MYSQL_HOST"),
        'user': os.getenv("MYSQL_USER"),
        'password': os.getenv("MYSQL_PASSWORD"),
        'database': os.getenv("MYSQL_DATABASE"),
        'port': 3306,
    }
    # 连接到MySQL数据库
    connection = mysql.connector.connect(**db_config)

    try:
        # 存储结果的字符串
        ddl = ""
        # 获取游标
        cursor = connection.cursor(buffered=True)
        # 遍历表信息
        for name in name_list:
            # 查询表信息
            cursor.execute(f"SHOW CREATE TABLE {name}")
            # 获取查询结果
            result = cursor.fetchone()
            if result:
                table_ddl = result[1] + ";\n"
                ddl += f"{table_ddl}\n"

        return ddl

    finally:
        # 关闭游标和连接
        cursor.close()
        connection.close()


def ask_database(query):
    print("=====数据库请求=====", query)
    global cursor
    db_config = {
        'host': os.getenv("MYSQL_HOST"),
        'user': os.getenv("MYSQL_USER"),
        'password': os.getenv("MYSQL_PASSWORD"),
        'database': os.getenv("MYSQL_DATABASE"),
        'port': 3306,
    }
    # 连接到MySQL数据库
    connection = mysql.connector.connect(**db_config)
    # 获取游标
    cursor = connection.cursor(buffered=True)
    cursor.execute(query)
    records = cursor.fetchall()
    return records


prompt = "统计每月每件商品的销售额"
# prompt = "这星期消费最高的用户是谁？他买了哪些商品？ 每件商品买了几件？花费多少？"
# prompt = "统计指定学员每个科目的错题总数量（包括：小测和练习），做错和半对的答案都算是错题。表的别名要简写"
prompt = ("""
根据如下数据表名称：\n\n
tb_course\n
stu_chapter_exercises_answer\n
stu_homework_exercises_answer\n
tb_chapter_exercises_homework\n
tb_chapter_exercises_explain_video\n\n
统计每个学员每个科目的错题总数量（包括：小测和练习），做错和半对的答案都算是错题。表的别名要简写。
返回结果包括如下列：\n\n
学员（student_id）、学科(subject)、错题数量(wrong_count)
"""
          )
messages = [
    {"role": "system", "content": "你是一个 MySQL 助手，可以帮助使用者编写 SQL 查询语句。"},
    {"role": "user", "content": prompt}
]
# response = get_sql_completion(messages, "gpt-4")

response = get_sql_completion(messages)
messages.append(response)  # 把大模型的回复加入到对话中


while 'tool_calls' in response and response['tool_calls'] is not None:
    # 1106 版新模型支持一次返回多个函数调用请求，所以要考虑到这种情况
    for tool_call in response['tool_calls']:
        args = json.loads(tool_call['function']['arguments'])
        if tool_call['function']['name'] == "get_table_ddl":
            result = get_table_ddl(**args)
        elif tool_call['function']['name'] == "ask_database":
            result = ask_database(**args)

        messages.append({
            "tool_call_id": tool_call['id'],  # 用于标识函数调用的 ID
            "role": "tool",
            "name": tool_call['function']['name'],
            "content": str(result)  # 数值result 必须转成字符串
        })

    response = get_sql_completion(messages)
    messages.append(response)  # 把大模型的回复加入到对话中




# 调用方法并打印结果
# database_schema_string = get_table_ddl(["stu_chapter_exercises_answer", "stu_homework_exercises_answer", "tb_chapter_exercises_homework", "tb_chapter_exercises_explain_video"])
# print(result)